package com.yth.utils;

import com.jayway.jsonpath.*;
import com.jayway.jsonpath.internal.EvaluationContext;
import com.jayway.jsonpath.internal.Path;
import com.jayway.jsonpath.internal.PathRef;
import com.jayway.jsonpath.internal.filter.*;
import com.jayway.jsonpath.internal.path.*;
import com.jayway.jsonpath.spi.json.JacksonJsonProvider;
import com.jayway.jsonpath.spi.json.JsonProvider;
import lombok.EqualsAndHashCode;
import lombok.Getter;
import lombok.ToString;

import java.util.*;

/**
 * @author yutianhong
 * @version 1.0
 * @since 2025/5/15 15:04
 */
public class JsonPathUtils {

    private static final Configuration JSONPATH_CONFIG = initializeConfiguration();

    /**
     * Reads values from an object using a JsonPath expression
     */
    public static List<Object> read(Object json, String path) {
        return JsonPath.using(JSONPATH_CONFIG).parse(json).read(path);
    }

    /**
     * Sets a value in an object at the location specified by a JsonPath expression
     */
    public static void set(Object json, String path, Object newValue) {
        JsonPath.using(JSONPATH_CONFIG).parse(json).set(path, newValue);
    }

    /**
     * Validates that a JsonPath expression returns exactly one result and returns it
     */
    public static Object readSingleResult(Object json, String path) {
        List<Object> results = read(json, path);
        if (results.isEmpty()) {
            return null;
        } else if (results.size() == 1) {
            return results.get(0);
        } else {
            throw new RuntimeException("JsonPath expected one result but found " + results.size() + " elements at path: " + path);
        }
    }

    /**
     * Sets a value at the specified JsonPath, creating path structure if it doesn't exist.
     * This method can handle complex JsonPath expressions with filter predicates.
     */
    public static void ensurePathAndSet(Object json, String path, Object newValue) {
        if (trySet(json, path, newValue)) {
            return;
        }

        CompiledPath compiledPath = compilePath(path);
        if (compiledPath.isFunctionPath()) {
            throw new IllegalArgumentException("Failed to set via JSONPath: path contains function. path: " + path);
        }

        // Process the path tokens to create missing structure
        boolean missingStructure = false;
        for (PathToken pathToken = compiledPath.getRoot(); pathToken.getNext() != null; pathToken = pathToken.getNext()) {
            PathToken next = pathToken.getNext();
            // Handles a predicate token in the path by creating necessary objects if needed.
            if (next instanceof PredicatePathToken || next instanceof PropertyPathToken) {
                PathToken nextNext = next.getNext();
                if (!missingStructure) {
                    next.setNext(null);
                    List<Object> readData = readByPath(compiledPath, json);
                    boolean pathExist = readData.stream().anyMatch(Objects::nonNull);
                    next.setNext(nextNext);   // restore
                    if (pathExist) {
                        // The path exists, continue processing
                        continue;
                    }
                    if (readData.size() > 1) {
                        throw new RuntimeException("Failed to set via JSONPath: not support for multiple values in path. path: " + path);
                    }
                    missingStructure = true;
                }

                if (next instanceof PredicatePathToken) {
                    handlePredicateToken(json, path, (PredicatePathToken) next, pathToken, compiledPath);

                    // Verify our operation worked
                    next.setNext(null);
                    boolean pathNotFound = readPathList(compiledPath, json).isEmpty();
                    next.setNext(nextNext);   // restore
                    if (pathNotFound) {
                        // Fast fail
                        throw new RuntimeException("Failed to set via JSONPath: path still not found after creating element. path: " + path);
                    }
                } else {
                    handlePropertyToken(json, path, (PropertyPathToken) next, compiledPath);
                }
            }
        }

        // Final attempt to set the value after ensuring path exists
        if (!trySet(json, path, newValue)) {
            throw new RuntimeException("Failed to set via JSONPath after creating path structure: path not found. path: " + path);
        }
    }

    private static void handlePropertyToken(Object json, String path, PropertyPathToken curr, CompiledPath compiledPath) {
        PathToken next = curr.getNext();
        if (next == null) {
            return; // no need to add property due to DEFAULT_PATH_LEAF_TO_NULL
        }
        if (next instanceof ArrayPathToken) {
            throw new RuntimeException("Failed to set via JSONPath: not support for array path in a not found value. path: " + path);
        }
        List<String> properties = ReflectUtils.getDeclaredFieldValue(PropertyPathToken.class, "properties", curr);
        if (properties == null || properties.isEmpty()) {
            throw new IllegalArgumentException("Failed to set via JSONPath: no properties found in path: " + path);
        }
        if (properties.size() > 1) {
            throw new IllegalArgumentException("Failed to set via JSONPath: not support for multiple properties in a single filter. path: " + path);
        }
        Object newValue;
        if (next instanceof PredicatePathToken) {
            newValue = JSONPATH_CONFIG.jsonProvider().createArray();
        } else {
            newValue = JSONPATH_CONFIG.jsonProvider().createMap();
        }
        curr.setNext(null);
        setByPath(compiledPath, json, newValue);
        curr.setNext(next);     // restore
    }

    private static void handlePredicateToken(Object json, String path, PredicatePathToken curr, PathToken prev, CompiledPath compiledPath) {
        Collection<Predicate> predicates = ReflectUtils.getDeclaredFieldValue(PredicatePathToken.class, "predicates", curr);
        if (predicates == null || predicates.isEmpty()) {
            throw new IllegalArgumentException("Failed to set via JSONPath: no predicates found in path: " + path);
        }
        if (predicates.size() > 1) {
            throw new IllegalArgumentException("Failed to set via JSONPath: not support for multiple predicates in a single filter. path: " + path);
        }
        Map<Path, Object> conditionMap = new HashMap<>();
        Deque<Pair<Predicate, Boolean>> stack = new ArrayDeque<>();
        Predicate compiledFilter = predicates.iterator().next();
        Predicate predicateRoot = ReflectUtils.getDeclaredFieldValue(compiledFilter.getClass(), "predicate", compiledFilter);   // risky reflect
        stack.push(Pair.of(predicateRoot, Boolean.TRUE));
        while (!stack.isEmpty()) {
            Pair<Predicate, Boolean> pair = stack.pop();
            Predicate predicate = pair.getLeft();
            Boolean isPositiveCondition = pair.getRight();
            if (predicate instanceof RelationalExpressionNode) {
                RelationalExpressionNode relationalExpressionNode = (RelationalExpressionNode) predicate;
                extractConditionFromRelationalNode(relationalExpressionNode, conditionMap, isPositiveCondition);
            } else if (predicate instanceof LogicalExpressionNode) {
                LogicalExpressionNode logicalExpressionNode = (LogicalExpressionNode) predicate;
                List<ExpressionNode> chain = ReflectUtils.getDeclaredFieldValue(LogicalExpressionNode.class, "chain", logicalExpressionNode);
                switch (logicalExpressionNode.getOperator()) {
                    case AND:
                    case OR:
                        chain.forEach(e -> stack.push(Pair.of(e, isPositiveCondition)));
                        break;
                    case NOT:
                        chain.forEach(e -> stack.push(Pair.of(e, !isPositiveCondition)));
                        break;
                    default:
                        throw new IllegalArgumentException("Failed to set via JSONPath: not support for logical operator. operator: " + logicalExpressionNode.getOperator() + ", path: " + path);
                }
            }
            // Other expression types are not handled for object creation
        }
        // Create new element with the extracted conditionsObject
        Object newElement = JSONPATH_CONFIG.jsonProvider().createMap();
        conditionMap.forEach((k, v) -> setByPath(k, newElement, v));
        // Find the array to add the new element
        prev.setNext(null);
        List<Object> readData = readByPath(compiledPath, json);
        Object target = readData.get(0);
        if (target instanceof Map) {
            // Set map
            setByPath(compiledPath, json, newElement);
        } else {
            // Add element to array
            addByPath(compiledPath, json, newElement);
        }
        prev.setNext(curr);    // restore
    }

    private static void extractConditionFromRelationalNode(RelationalExpressionNode relationalExpressionNode, Map<Path, Object> conditionMap, boolean isPositiveCondition) {
        RelationalOperator relationalOperator = ReflectUtils.getDeclaredFieldValue(RelationalExpressionNode.class, "relationalOperator", relationalExpressionNode);
        if (isPositiveCondition && relationalOperator != RelationalOperator.EQ && relationalOperator != RelationalOperator.TSEQ) {
            return;
        }
        if (!isPositiveCondition && relationalOperator != RelationalOperator.NE && relationalOperator != RelationalOperator.TSNE) {
            return;
        }
        ValueNode left = ReflectUtils.getDeclaredFieldValue(RelationalExpressionNode.class, "left", relationalExpressionNode);
        if (!left.isPathNode()) {
            return; // Must be a path on the left side
        }
        Path leftPath = left.asPathNode().getPath();
        ValueNode right = ReflectUtils.getDeclaredFieldValue(RelationalExpressionNode.class, "right", relationalExpressionNode);
        Object value;
        if (right.isBooleanNode()) {
            value = right.asBooleanNode().getBoolean();
        } else if (right.isNumberNode()) {
            value = right.asNumberNode().getNumber();
        } else if (right.isStringNode()) {
            value = right.asStringNode().getString();
        } else if (right.isNullNode()) {
            value = null;
        } else {
            return; // Unsupported value type
        }
        conditionMap.put(leftPath, value);
    }

    private static <T> T readByPath(CompiledPath path, Object json) {
        EvaluationContext evaluationContext = path.evaluate(json, json, JSONPATH_CONFIG);
        if (evaluationContext.getPathList().isEmpty()) {
            return (T) JSONPATH_CONFIG.jsonProvider().createArray();
        }
        Object res = ReflectUtils.getDeclaredFieldValue(EvaluationContextImpl.class, "valueResult", evaluationContext);
        return (T) res;
    }

    private static List<String> readPathList(Path path, Object jsonObject) {
        EvaluationContext evaluationContext = path.evaluate(jsonObject, jsonObject, JSONPATH_CONFIG);
        if (evaluationContext.getPathList().isEmpty()) {
            return (List<String>) JSONPATH_CONFIG.jsonProvider().createArray();
        }
        return evaluationContext.getPath();
    }

    private static void setByPath(Path path, Object jsonObject, Object newVal) {
        EvaluationContext evaluationContext = path.evaluate(jsonObject, jsonObject, JSONPATH_CONFIG, true);
        if (evaluationContext.getPathList().isEmpty()) {
            throw new RuntimeException("Failed to set via JSONPath: path not found. path: " + path);
        }
        for (PathRef updateOperation : evaluationContext.updateOperations()) {
            updateOperation.set(newVal, JSONPATH_CONFIG);
        }
    }

    private static void addByPath(Path path, Object jsonObject, Object value) {
        EvaluationContext evaluationContext = path.evaluate(jsonObject, jsonObject, JSONPATH_CONFIG, true);
        if (evaluationContext.getPathList().isEmpty()) {
            throw new RuntimeException("Failed to add via JSONPath: path not found. path: " + path);
        }
        for (PathRef updateOperation : evaluationContext.updateOperations()) {
            updateOperation.add(value, JsonPathUtils.JSONPATH_CONFIG);
        }
    }

    private static boolean trySet(Object json, String path, Object newValue) {
        List<Object> results = JsonPath.using(JSONPATH_CONFIG.addOptions(Option.AS_PATH_LIST)).parse(json).read(path);
        if (results == null || results.isEmpty()) {
            return false;
        }
        JsonPath.using(JSONPATH_CONFIG).parse(json).set(path, newValue);
        return true;
    }

    private static CompiledPath compilePath(String path) {
        // Can't cache path due to we need to call setNext(null) during process
        Path compiled = PathCompiler.compile(path);
        if (compiled instanceof CompiledPath) {
            return (CompiledPath) compiled;
        } else {
            throw new RuntimeException("JsonPath compiled is not a CompiledPath. path: " + path);
        }
    }

    private static Configuration initializeConfiguration() {
        JsonProvider customJacksonJsonProvider = new JacksonJsonProvider() {
            @Override
            public void setProperty(Object obj, Object key, Object value) {
                if (isMap(obj)) {
                    ((Map<Object, Object>) obj).put(key.toString(), value);
                } else if (isArray(obj)) {
                    List<Object> array = (List<Object>) (obj);
                    int index;
                    if (key != null) {
                        index = key instanceof Integer ? (Integer) key : Integer.parseInt(key.toString());
                    } else {
                        index = array.size();
                    }

                    if (index == array.size()) {
                        array.add(value);
                    } else {
                        array.set(index, value);
                    }
                } else {
                    throw new JsonPathException("setProperty operation cannot be used with " + (obj != null ? obj.getClass().getName() : "null"));
                }
            }
        };
        return Configuration.defaultConfiguration()
                .addOptions(Option.SUPPRESS_EXCEPTIONS, Option.ALWAYS_RETURN_LIST, Option.DEFAULT_PATH_LEAF_TO_NULL)
                .jsonProvider(customJacksonJsonProvider);
    }

    @Getter
    @ToString
    @EqualsAndHashCode
    private static class Pair<L, R> {
        private final L left;
        private final R right;

        private Pair(L left, R right) {
            this.left = left;
            this.right = right;
        }

        public static <L, R> Pair<L, R> of(L left, R right) {
            return new Pair<>(left, right);
        }
    }

    private JsonPathUtils() {
    }
}
